import os
import cv2
import numpy as np
from PIL import Image
from ultralytics import YOLO
import face_recognition

class FaceBatchProcessor:
    def __init__(self, input_folder, output_folder, model_path='/home/jaesonzhang/working/photo_person/yolov9m-face-lindevs.pt'):
        self.input_folder = input_folder
        self.output_folder = output_folder
        self.face_encodings = []  # 存储特征向量
        self.face_info = []       # 存储人脸信息
        self.model = YOLO(model_path)  # 加载YOLOv8-face模型[1,6](@ref)

    def process_batch_photos(self):
        """批量处理照片"""
        photo_files = [f for f in os.listdir(self.input_folder) 
                      if f.lower().endswith(('.jpg', '.jpeg', '.png'))]
        
        for i, photo_file in enumerate(photo_files):
            print(f"处理照片 {i+1}/{len(photo_files)}: {photo_file}")
            self.process_single_photo(photo_file)
            
    def process_single_photo(self, photo_file):
        """处理单张照片"""
        image_path = os.path.join(self.input_folder, photo_file)
        
        # 使用OpenCV加载图片（YOLOv8兼容格式）
        image_cv = cv2.imread(image_path)
        image_rgb = cv2.cvtColor(image_cv, cv2.COLOR_BGR2RGB)  # 转换为RGB格式
        
        # YOLOv8人脸检测
        results = self.model.predict(image_rgb, conf=0.5)  # 置信度阈值设为0.5[1](@ref)
        boxes = results[0].boxes.xyxy.cpu().numpy().astype(int)  # 获取检测框坐标[x1,y1,x2,y2]
        
        print(f"在 {photo_file} 中检测到 {len(boxes)} 张人脸")
        
        # 转换为face_recognition所需的格式(top, right, bottom, left)
        face_locations = [(box[1], box[2], box[3], box[0]) for box in boxes]
        
        # 计算人脸特征（使用原图RGB格式）
        face_encodings = face_recognition.face_encodings(image_rgb, face_locations)
        
        # 处理每个检测到的人脸
        for j, (face_location, face_encoding) in enumerate(zip(face_locations, face_encodings)):
            self.save_and_record_face(face_encoding, face_location, image_rgb, photo_file, j)
    
    def save_and_record_face(self, encoding, location, image, original_file, face_index):
        """保存人脸图片并记录特征"""
        # 提取人脸区域（注意YOLOv8的location顺序与face_recognition一致）
        top, right, bottom, left = location
        face_image = image[top:bottom, left:right]
        
        # 生成唯一文件名
        base_name = os.path.splitext(original_file)[0]
        face_filename = f"{base_name}_face{face_index}.jpg"
        face_path = os.path.join(self.output_folder, "faces", face_filename)
        
        # 保存人脸图片（PIL格式）
        os.makedirs(os.path.dirname(face_path), exist_ok=True)
        pil_image = Image.fromarray(face_image)
        pil_image.save(face_path)
        
        # 记录信息
        self.face_encodings.append(encoding)
        self.face_info.append({
            'original_file': original_file,
            'face_file': face_filename,
            'face_index': face_index,
            'location': location
        })

        
        
from sklearn.cluster import DBSCAN
import matplotlib.pyplot as plt
from collections import defaultdict

class FaceClassifier:
    def __init__(self, face_encodings, face_info):
        self.face_encodings = face_encodings
        self.face_info = face_info
        self.labels = None
        
    def cluster_faces(self, eps=0.5, min_samples=2):
        """使用DBSCAN进行人脸聚类"""
        # 转换为numpy数组
        X = np.array(self.face_encodings)
        
        # 聚类
        clustering = DBSCAN(eps=eps, min_samples=min_samples, metric='euclidean').fit(X)
        self.labels = clustering.labels_
        
        # 统计聚类结果
        n_clusters = len(set(self.labels)) - (1 if -1 in self.labels else 0)
        n_noise = list(self.labels).count(-1)
        
        print(f'聚类完成: 发现 {n_clusters} 个人物, {n_noise} 个无法分类的人脸')
        return self.labels
    
    def organize_by_person(self):
        """按人物整理分类结果"""
        if self.labels is None:
            raise ValueError("请先执行聚类操作")
            
        person_groups = defaultdict(list)
        
        for label, info in zip(self.labels, self.face_info):
            person_id = f"person_{label}" if label != -1 else "unknown"
            person_groups[person_id].append(info)
            
        return person_groups
    
    def generate_classification_report(self, output_folder):
        """生成分类报告和整理文件"""
        person_groups = self.organize_by_person()
        
        # 为每个人物创建文件夹
        for person_id, faces in person_groups.items():
            person_folder = os.path.join(output_folder, "classified", person_id)
            os.makedirs(person_folder, exist_ok=True)
            
            # 复制人脸图片到对应文件夹
            for face_info in faces:
                src_path = os.path.join(output_folder, "faces", face_info['face_file'])
                dst_path = os.path.join(person_folder, face_info['face_file'])
                
                if os.path.exists(src_path):
                    # 复制文件
                    import shutil
                    shutil.copy2(src_path, dst_path)
        
        # 生成统计报告
        report_path = os.path.join(output_folder, "classification_report.txt")
        with open(report_path, 'w', encoding='utf-8') as f:
            f.write("人脸分类报告\n")
            f.write("=" * 50 + "\n")
            for person_id, faces in person_groups.items():
                f.write(f"\n{person_id}: {len(faces)} 张人脸\n")
                original_files = set(info['original_file'] for info in faces)
                f.write(f"  来源照片: {', '.join(original_files)}\n")
                
def main():
    # 配置路径
    input_folder = "input_images"      # 原始照片文件夹
    output_folder = "output_results"   # 输出结果文件夹
    
    # 1. 人脸检测与提取
    processor = FaceBatchProcessor(input_folder, output_folder)
    processor.process_batch_photos()
    
    # 2. 人脸聚类分类
    classifier = FaceClassifier(processor.face_encodings, processor.face_info)
    labels = classifier.cluster_faces(eps=0.5, min_samples=2)
    
    # 3. 生成分类结果
    classifier.generate_classification_report(output_folder)
    
    print("批量人脸分类完成！")

if __name__ == "__main__":
    main()
